# -*- coding: utf-8 -*-

import os
import shutil
import tensorflow as tf
from skimage import io

flags = tf.app.flags
flags.DEFINE_string( 'data_dir' , '' , "")
flags.DEFINE_string( 'data_info' , '' , "" )
flags.DEFINE_string( 'export_dir' , '' , "")

FLAGS = flags.FLAGS

class voc_xml( object ):
    headstr = """\
<annotation>
    <folder>John Hush</folder>
    <filename>%s</filename>
    <source>
        <database>My Database</database>
        <annotation>John Hush HAHAHAHA</annotation>
        <image>flickr</image>
        <flickrid>NULL</flickrid>
    </source>
    <owner>
        <flickrid>NULL</flickrid>
        <name>company</name>
    </owner>
    <size>
        <width>%d</width>
        <height>%d</height>
        <depth>%d</depth>
    </size>
    <segmented>0</segmented>
"""
    objstr = """\
    <object>
        <name>%s</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>%d</xmin>
            <ymin>%d</ymin>
            <xmax>%d</xmax>
            <ymax>%d</ymax>
        </bndbox>
    </object>
"""

    tailstr =\
    """
</annotation>
"""

def write_xml( f , head , bboxes ):
    # write the three parts into .xml files
    with open( f , 'w' ) as fw:
        fw.write( head )

        for bbox in bboxes:
            fw.write( voc_xml.objstr %( 'face' , bbox[0] , \
                    bbox[1] , bbox[2] , bbox[3] ) )

        fw.write( voc_xml.tailstr )

def dataset_generate( data_dir , data_info , export_dir ):
    #refresh export dir
    anno_dir = os.path.join( export_dir, 'Annotations' )
    jpeg_dir = os.path.join( export_dir, 'JPEGImages' )
    main_dir = os.path.join( export_dir , 'ImageSets/Main')

    if shutil.os.path.exists( anno_dir ):
        shutil.rmtree( anno_dir )
    if shutil.os.path.exists( jpeg_dir ):
        shutil.rmtree( jpeg_dir )
    if shutil.os.path.exists( os.path.join( export_dir, 'ImageSets' ) ):
        shutil.rmtree( os.path.join( export_dir, 'ImageSets' ) )

    shutil.os.mkdir( anno_dir )
    shutil.os.mkdir( jpeg_dir )
    shutil.os.makedirs( main_dir )

    data_indicator = "train.txt"
    img_info_txt = os.path.join( export_dir , 'ImageSets/Main', \
            data_indicator )

    with open( data_info , 'r' ) as fr1:
        face_detection_info_list = fr1.readlines()
        #face_detection_info_list = [ s for s in face_detection_info_list \
        #        if len(s.strip().split(' ')) == 5 ]

    with open( img_info_txt , 'a' ) as fr2:
        for face_detection_info in face_detection_info_list:
            info = face_detection_info.strip().split( ' ' )

            if len( info ) == 1 or info[1][-4:] == '.jpg':
                continue
            bbox_num = int( ( len( info ) -1 )/4 )

            print( "filename is: %s" % info[0] )

            img = io.imread( os.path.join( data_dir , info[0] ) )
            head = voc_xml.headstr %( info[0] , img.shape[1] , \
                    img.shape[0] , img.shape[2] )

            bboxes = []
            
            for bbox_index in range( bbox_num ):
                index_start = 1 + 4 * bbox_index
                bbox = [ int(x) for x in info[index_start : index_start+4] ]
                bboxes.append( bbox )

            xml_name = os.path.join( export_dir , "Annotations" , \
                    info[0][:-4] + '.xml' )
            
            #write Annotations
            write_xml( xml_name , head , bboxes )
            #write /Main things
            fr2.write( '%s\n' %(info[0]) )
            #copy img to /JPEGImage dir
            shutil.copyfile( os.path.join(data_dir, info[0]) , \
                    os.path.join( export_dir, 'JPEGImages' , info[0] ) )


if __name__ == "__main__":
    dataset_generate(FLAGS.data_dir , FLAGS.data_info, FLAGS.export_dir )
